"""
Case Type   : DATAVEC模块
Case Name   : vector字段创建hnsw l2距离索引，查询参数hnsw_ef_search与索引选项ef_construction参数覆盖验证
Create At   : 2025.4.18
Owner       : wang-tianjing1
Description :
    step1: 创建3张表
    step2: 创建自定义插入函数
    step3: 使用自定义插入函数向3张表插入相同数据
    step4: 创建索引，指定ef_construction，小于默认值64
    step5: 创建索引，不指定ef_construction
    step6: 创建索引，指定ef_construction，大于默认值64
    step7: 设置不同hnsw_ef_search参数，进行查询操作
    step8: 清理环境
Expect      :
    step1: 成功
    step2: 成功
    step3: 成功
    step4: 成功
    step5: 成功
    step6: 成功
    step7: 成功，查询结果相同
    step8: 成功
History     :
"""

import os
import unittest
import time
from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node


class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        self.log = Logger()
        self.log.info(f'----{os.path.basename(__file__)}:start----')
        self.sh = CommonSH('PrimaryDbUser')
        self.pri_dbuser = Node(node='PrimaryDbUser')
        self.sta_1_sh = CommonSH('Standby1DbUser')
        self.constant = Constant()
        self.com = Common()
        self.comsql = CommonSQL('PrimaryDbUser')
        self.tb_name_1 = 't_datavec_index0120_1'
        self.tb_name_2 = 't_datavec_index0120_2'
        self.tb_name_3 = 't_datavec_index0120_3'
        self.fun_insert = 'f_datavec_index0120'

    def test_main(self):
        step_txt = '----step1: 创建表; expect: 成功----'
        self.log.info(step_txt)
        create_tb = f"drop table if exists {self.tb_name_1},{self.tb_name_2},{self.tb_name_3}; " \
                    f"create table {self.tb_name_1} (id bigserial primary key, val vector(5)); " \
                    f"create table {self.tb_name_2} (id bigserial primary key, val vector(5)); " \
                    f"create table {self.tb_name_3} (id bigserial primary key, val vector(5));"
        result = self.sh.execut_db_sql(create_tb)
        self.log.info(result)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result, '执行失败:' + step_txt)

        step_txt = '----step2: 创建自定义插入函数; expect: 成功----'
        self.log.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        self.log.info(result)
        self.assertEqual(result, self.fun_insert, '执行失败: 创建自定义插入函数失败')

        step_txt = '----step3: 表插入相同数据; expect: 成功----'
        self.log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(5, 10000, '{self.tb_name_1}', 'val','vector');"
        result = self.sh.execut_db_sql(insert_sql)
        self.log.info(result)
        self.assertIn('SUCCESS', result, '执行失败: 使用自定义函数向表1插入数据失败')

        sql = f"insert into {self.tb_name_2} select * from {self.tb_name_1};"
        result = self.sh.execut_db_sql(sql)
        self.log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result, '执行失败: 向表2插入数据失败')

        sql = f"insert into {self.tb_name_3} select * from {self.tb_name_1};"
        result = self.sh.execut_db_sql(sql)
        self.log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result, '执行失败: 向表3插入数据失败')

        step_txt = '----step4: 创建索引，指定ef_construction，小于默认值64; expect: 成功----'
        self.log.info(step_txt)
        create_index_1 = f"create index idx_{self.tb_name_1} on {self.tb_name_1} using hnsw (val vector_l2_ops) with(m=5,ef_construction=32);"
        result = self.sh.execut_db_sql(create_index_1)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result, '执行失败: 创建索引1失败')

        step_txt = '----step5: 创建索引，不指定ef_construction; expect: 成功----'
        self.log.info(step_txt)
        create_index_2 = f"create index idx_{self.tb_name_2} on {self.tb_name_2} using hnsw (val vector_l2_ops) with(m=5);"
        result = self.sh.execut_db_sql(create_index_2)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result, '执行失败: 创建索引2失败')

        step_txt = '----step6: 创建索引，指定ef_construction，大于默认值64; expect: 成功----'
        self.log.info(step_txt)
        create_index_3 = f"create index idx_{self.tb_name_3} on {self.tb_name_3} using hnsw (val vector_l2_ops) with(m=5,ef_construction=128);"
        result = self.sh.execut_db_sql(create_index_3)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result, '执行失败: 创建索引3失败')

        step_txt = '----step7: 设置不同hnsw_ef_search参数，进行查询操作; expect: 功能正常----'
        self.log.info(step_txt)
        hnsw_ef_search_values = [10, 64, 300]
        query_vector = '[49,490.2,498,491,49]'
        for hnsw_ef_search in hnsw_ef_search_values:
            sql = f"set enable_indexscan=on; set enable_seqscan=off; set hnsw_ef_search={hnsw_ef_search};" \
                  f"select *, val<-> '{query_vector}' from {self.tb_name_1} order by val<-> '{query_vector}' limit 20;" \
                  f"select *, val<-> '{query_vector}' from {self.tb_name_2} order by val<-> '{query_vector}' limit 20;" \
                  f"select *, val<-> '{query_vector}' from {self.tb_name_3} order by val<-> '{query_vector}' limit 20;"
            result_all = self.sh.execut_db_sql(sql)

            # 去除 SET 语句的影响，提取每个查询的结果
            result_lines = result_all.split('\n')
            query_results = []
            current_result = []
            for line in result_lines:
                if line.strip() == '':
                    if current_result:
                        query_results.append('\n'.join(current_result))
                        current_result = []
                elif not line.startswith('SET'):
                    current_result.append(line)
            if current_result:
                query_results.append('\n'.join(current_result))

            result_1 = query_results[0]
            result_2 = query_results[1]
            result_3 = query_results[2]

            self.log.info(f"{self.tb_name_1}: {result_1}, {self.tb_name_2}: {result_2}, {self.tb_name_2}: {result_3}")
            self.assertEqual(result_1, result_2, f'执行失败: hnsw_ef_search={hnsw_ef_search} 时表1和表2返回值不同')
            self.assertEqual(result_1, result_3, f'执行失败: hnsw_ef_search={hnsw_ef_search} 时表1和表3返回值不同')

    def tearDown(self):
        step_txt = '----step8: 清理环境; expect:成功----'
        self.log.info(step_txt)
        drop_cmd = f"drop table if exists {self.tb_name_2},{self.tb_name_3};"
        drop_result1 = self.sh.execut_db_sql(drop_cmd)
        self.log.info(drop_result1)
        drop_cmd = f"drop function if exists {self.fun_insert};"
        drop_result2 = self.sh.execut_db_sql(drop_cmd)
        self.log.info(drop_result2)
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, drop_result1, '执行失败:' + step_txt)
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, drop_result2, '执行失败:' + step_txt)
        self.log.info(f'----{os.path.basename(__file__)}:end----')


